{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['/mnt/hdfs/hongli/dataset/hotpotqa/eval_100.json',\n",
       " '/mnt/hdfs/hongli/dataset/hotpotqa/eval_1600.json',\n",
       " '/mnt/hdfs/hongli/dataset/hotpotqa/eval_200.json',\n",
       " '/mnt/hdfs/hongli/dataset/hotpotqa/eval_3200.json',\n",
       " '/mnt/hdfs/hongli/dataset/hotpotqa/eval_400.json',\n",
       " '/mnt/hdfs/hongli/dataset/hotpotqa/eval_50.json',\n",
       " '/mnt/hdfs/hongli/dataset/hotpotqa/eval_6400.json',\n",
       " '/mnt/hdfs/hongli/dataset/hotpotqa/eval_800.json']"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import glob\n",
    "from transformers import AutoTokenizer\n",
    "files = glob.glob(\"/mnt/hdfs/hongli/dataset/hotpotqa/eval*.json\")\n",
    "tok = AutoTokenizer.from_pretrained(\"/mnt/hdfs/hongli/model/Qwen2.5-0.5B-Instruct\")\n",
    "files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/mnt/hdfs/hongli/dataset/hotpotqa/eval_100.json max 15644 min 12351 avg 14029.765625\n",
      "/mnt/hdfs/hongli/dataset/hotpotqa/eval_1600.json max 231703 min 217296 avg 225270.6796875\n",
      "/mnt/hdfs/hongli/dataset/hotpotqa/eval_200.json max 30832 min 25546 avg 28069.4375\n",
      "/mnt/hdfs/hongli/dataset/hotpotqa/eval_3200.json max 462065 min 438033 avg 451415.75\n",
      "/mnt/hdfs/hongli/dataset/hotpotqa/eval_400.json max 60824 min 51695 avg 56152.8671875\n",
      "/mnt/hdfs/hongli/dataset/hotpotqa/eval_50.json max 8420 min 5821 avg 6999.5546875\n",
      "/mnt/hdfs/hongli/dataset/hotpotqa/eval_6400.json max 915357 min 890081 avg 903061.984375\n",
      "/mnt/hdfs/hongli/dataset/hotpotqa/eval_800.json max 117305 min 108084 avg 112227.21875\n"
     ]
    }
   ],
   "source": [
    "import json\n",
    "for fn in files:\n",
    "    with open(fn) as f:\n",
    "        data = json.load(f)\n",
    "    tokens  = [len(x) for x in tok([j['context'] for j in data])['input_ids']]\n",
    "    print(fn, \"max\", max(tokens), \"min\", min(tokens), \"avg\", sum(tokens) / len(tokens))"
   ]
  }
 ],
 "metadata": {
  "fileId": "3c8c6878-99ec-475c-8f3b-996b00524ea1",
  "filePath": "/opt/tiger/verl/taskutils/memory_data/token_count.ipynb",
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
